package com.adam.example.aggregation.bucket;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.DistanceUnit;
import co.elastic.clients.elasticsearch._types.aggregations.*;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import com.adam.example.ElasticsearchClientHolder;
import com.adam.example.util.JsonUtil;

import java.io.IOException;
import java.util.Arrays;
import java.util.List;

/**
 * 执行报错：无法解析ranges字段。原因是范围值为字符串类型。
 */
public class GeoDistanceAggregationExample {

    public static void main(String[] args) throws IOException {
        ElasticsearchClient elasticsearchClient = ElasticsearchClientHolder.getINSTANCE().initializeAndGetClient();

        SearchRequest searchRequest = SearchRequest.of(builder -> builder
                        .index("test_geo_point")
                        .size(0)
                        .aggregations("agg", aggregationBuilder -> aggregationBuilder
                                .geoDistance(geoDistanceAggregationBuilder -> geoDistanceAggregationBuilder
                                        .field("location")
                                        .origin(geoLocationBuilder -> geoLocationBuilder
                                                .text("34.3412700000,108.9398400000")
                                        )
                                        .unit(DistanceUnit.Kilometers)
                                        .ranges(Arrays.asList(
                                                AggregationRange.of(aggregationRangeBuilder -> aggregationRangeBuilder
                                                        .to("500")
                                                ),
                                                AggregationRange.of(aggregationRangeBuilder -> aggregationRangeBuilder
                                                        .from("500").to("1000")
                                                ),
                                                AggregationRange.of(aggregationRangeBuilder -> aggregationRangeBuilder
                                                        .from("1000")
                                                )
                                        ))
                                )
                        )
        );
        System.out.println(JsonUtil.searchRequestToJson(searchRequest, elasticsearchClient));
        SearchResponse<Void> searchResponse = elasticsearchClient.search(searchRequest, Void.class);
        Aggregate aggregate = searchResponse.aggregations().get("agg");
        GeoDistanceAggregate geoDistanceAggregate = aggregate.geoDistance();
        Buckets<RangeBucket> buckets = geoDistanceAggregate.buckets();
        List<RangeBucket> bucketList = buckets.array();
        for(RangeBucket rangeBucket: bucketList) {
            System.out.println(rangeBucket.key() + ":" + rangeBucket.docCount());
        }

        ElasticsearchClientHolder.getINSTANCE().close();
    }

}
